from itertools import count
from typing import Tuple

from lagent.schema import AgentMessage


class Hook:
    """
    钩子基类
    
    提供代理和动作执行器的生命周期钩子方法
    允许在代理处理和动作执行的不同阶段插入自定义逻辑
    """

    def before_agent(
        self,
        agent,
        message: Tuple[AgentMessage],
        session_id: int,
    ):
        """
        代理处理消息前的钩子方法
        
        参数:
            agent: 代理实例
            message: 待处理的消息元组
            session_id: 会话ID
            
        返回:
            None 或修改后的消息
        """
        pass

    def after_agent(
        self,
        agent,
        message: AgentMessage,
        session_id: int,
    ):
        """
        代理处理消息后的钩子方法
        
        参数:
            agent: 代理实例
            message: 处理后的消息
            session_id: 会话ID
            
        返回:
            None 或修改后的消息
        """
        pass

    def before_action(
        self,
        executor,
        message: AgentMessage,
        session_id: int,
    ):
        """
        动作执行前的钩子方法
        
        参数:
            executor: 动作执行器实例
            message: 待执行的动作消息
            session_id: 会话ID
            
        返回:
            None 或修改后的消息
        """
        pass

    def after_action(
        self,
        executor,
        message: AgentMessage,
        session_id: int,
    ):
        """
        动作执行后的钩子方法
        
        参数:
            executor: 动作执行器实例
            message: 执行结果消息
            session_id: 会话ID
            
        返回:
            None 或修改后的消息
        """
        pass


class RemovableHandle:
    """
    可移除的钩子句柄类
    
    用于管理钩子的生命周期，支持动态移除钩子
    
    属性:
        _id_iter: 类级别的ID生成器
        hooks_dict: 存储钩子的字典
        id: 当前钩子的唯一标识符
    """
    _id_iter = count(0)  # 从0开始的自增ID生成器

    def __init__(self, hooks_dict):
        """
        初始化钩子句柄
        
        参数:
            hooks_dict: 存储钩子的字典
        """
        self.hooks_dict = hooks_dict
        self.id = next(self._id_iter)  # 生成唯一ID

    def remove(self):
        """
        从钩子字典中移除当前钩子
        """
        del self.hooks_dict[self.id]
